package com.gitee.wsl.platform.thread.task

import com.gitee.wsl.platform.thread.task.api.Fiber
import com.gitee.wsl.platform.thread.task.bean.Task
import org.jetbrains.annotations.ApiStatus
import org.jetbrains.annotations.Blocking
import java.util.concurrent.ExecutionException
import java.util.concurrent.Executor
import java.util.function.Supplier


@ApiStatus.Internal
internal interface TaskExecutor : Executor {

    fun resumeOnExecutor(runnable: Runnable)

    companion object {
        fun from(executor: Executor): TaskExecutor {
            return executor as? TaskExecutor ?: TaskExecutorWithForkedResume(executor)
        }
    }
}

@ApiStatus.Internal
internal object TaskLocalContext {
    fun signalTheStartOfBlockingCall() {
        // clears the trampoline first
        val executor = localExecutor.get()
        Trampoline.forkAll(
            executor ?: TaskExecutors.sharedBlockingIO()
        )
    }

    fun isCurrentExecutor(executor: TaskExecutor?): Boolean {
        val currentExecutor = localExecutor.get()
        return currentExecutor === executor
    }

    fun getAndSetExecutor(executor: TaskExecutor?): TaskExecutor? {
        val oldExecutor = localExecutor.get()
        localExecutor.set(executor)
        return oldExecutor
    }

    private val localExecutor: ThreadLocal<TaskExecutor?> = ThreadLocal.withInitial<TaskExecutor?>(Supplier { null })
}

@ApiStatus.Internal
internal class TaskExecutorWithForkedResume(private val executor: Executor) : TaskExecutor {
    override fun execute(command: Runnable) {
        executor.execute(Runnable {
            val oldExecutor = TaskLocalContext.getAndSetExecutor(this)
            try {
                command.run()
            } finally {
                TaskLocalContext.getAndSetExecutor(oldExecutor)
            }
        })
    }

    override fun resumeOnExecutor(runnable: Runnable) {
        if (TaskLocalContext.isCurrentExecutor(this)) {
            Trampoline.execute(runnable)
        } else {
            execute(runnable)
        }
    }
}

@Blocking
@Throws(InterruptedException::class, ExecutionException::class)
fun <T> Executor?.runBlockingUninterruptible(
    task: Task<T>,
): T? {
    val executor = this

    val fiber = task.runFiber(executor)

    fiber.joinBlockingUninterruptible()
    try {
        return fiber.getResultOrThrow()
    } catch (e: TaskCancellationException) {
        throw ExecutionException(e)
    } catch (e: Fiber.NotCompletedException) {
        throw ExecutionException(e)
    }
}